#include <omp.h>
#include <malloc.h>

#define THREADS (8)
#define BLOCK_SIZE (64) //128

void MatrixMultiplication(double* A, double* B, double *C, int aH, int aW, int bW)
{   
	int block_size_i = BLOCK_SIZE;
	while (aH % block_size_i != 0)
		block_size_i /= 2;

	int block_size_j = BLOCK_SIZE;
	while (bW % block_size_j != 0)
		block_size_j /= 2;

	int num_blocks_i = aH / block_size_i;
	int num_blocks_j = bW / block_size_j;
	omp_set_num_threads(THREADS);
	double* rows[THREADS];
	for (int i=0; i<THREADS; i++)
	{
		rows[i] = (double*) _mm_malloc(aW*block_size_i*sizeof(double), 32);
	}
	#pragma omp parallel for
	for (int ib = 0; ib < num_blocks_i; ib++)
	{
		double* row = rows[omp_get_thread_num()];
		for (int i=0; i<block_size_i; i++)
		{
			
			int i_real = i + ib*block_size_i;
			#pragma simd
			for (int j=0; j<aW; j++)
			{
				row[i*aW + j] = A[i_real + j*aH];
			}
		}
		for (int jb = 0; jb < num_blocks_j; jb++)
		{
			for (int i=0; i<block_size_i; i++)
			{
				int i_real = i + ib*block_size_i;
				#pragma vector nontemporal(C)
				for (int j=jb*block_size_j; j<(jb+1)*block_size_j; j++)
				{
					double sum = 0;
					#pragma simd
					for (int m=0; m<aW; m++)
					{
						sum += row[m + i*aW] * B[j*aW + m];
					}
					C[i_real + j*aH] = sum;
				}
			}
		}
	}
}

